# -*- coding: UTF-8 -*- 
import json

# "ori_index" : int
# "imagenet_label": str
# "imagenet_string": str
# "cn_string": str
all_cate = {
    '1':{"desc":"人物",
          "content":[]
	},
	'2':{"desc":"动物",
          "content":[]
	},
	'3':{"desc":"植物",
          "content":[]
	},
	'4':{"desc":"自然风景",
          "content":[]
	},
	'5':{"desc":"人造物",
          "content":[]
	},
    '6':{"desc":"建筑",
          "content":[]
    },
    '8':{"desc":"丢弃",
          "content":[]
    }
}

all_class_f = open('imagenet_class.txt', 'r', encoding='utf-8').readlines()

all_info_map = {}


for i, item in enumerate(all_class_f):
    index = item.find(',')
    info = item[:index].split()
    imagenet_string = item[index+1:].strip()
    cn_name = info[2]

    if cn_name not in all_info_map:
        all_info_map[cn_name] = []

    all_info_map[cn_name].append(info+[imagenet_string])


num = 0
for key, item in all_info_map.items():
    print(f'类别=={key}== len:{len(item)} {num}/1000')

    c = input()
    while(not c.isdigit() or int(c) not in [1, 2, 3, 4, 5, 6, 8, 9]):
        print('1:人物  2:动物  3:植物 4:自然风景 5:人造物品 6:建筑 8:丢弃  9:退出')
        c = input()

    if(int(c) == 9): break

    cate = all_cate[c]
    for ii in item:
        cate_i = {'ori_index':int(ii[0]), 
                  'imagenet_label':ii[1], 
                  'imagenet_string':ii[3], 
                  'cn_string':ii[2]}
        cate['content'].append(cate_i)
    num += len(item)

# print(len(all_info_map))

    # print(f'no.{i}: {info[-1]}  {imagenet_string}')
    
    # c = input()
    # while(not c.isdigit() or int(c) not in [1, 2, 3, 4, 5, 8, 9]):
    #     print('1:人物  2:动物  3:植物 4:自然风景 5:人造物品  8:丢弃  9:退出')
    #     c = input()
    # if(int(c) == 9): break

    # cate = all_cate[c]
    # cate_i = {'ori_index':int(info[0]), 
    #           'imagenet_label':info[1], 
    #           'imagenet_string':imagenet_string, 
    #           'cn_string':info[2]}
    # cate['content'].append(cate_i)

of = open('category_info_gen.json', 'w', encoding='utf-8')


dmp = [item for k, item in all_cate.items()]
of.write(json.dumps(dmp, ensure_ascii=False))
of.close()